from setuptools import setup, find_packages

# 主要依赖
MAIN_REQUIREMENTS = [
    "accelerate>=0.21.0",
    "datasets>=3.1.0",
    "optimum>=1.23.3",
    "peft>=0.13.2",
    "transformers>=4.36.0",
    "wandb==0.18.7",
]

# CUDA相关依赖
CUDA_REQUIREMENTS = [
    "nvidia-cublas-cu12==12.4.5.8",
    "nvidia-cuda-cupti-cu12==12.4.127",
    "nvidia-cuda-nvrtc-cu12==12.4.127",
    "nvidia-cuda-runtime-cu12==12.4.127",
    "nvidia-cudnn-cu12==9.1.0.70",
    "nvidia-cufft-cu12==11.2.1.3",
    "nvidia-curand-cu12==10.3.5.147",
    "nvidia-cusolver-cu12==11.6.1.9",
    "nvidia-cusparse-cu12==12.3.1.170",
    "nvidia-nccl-cu12==2.21.5",
    "nvidia-nvtx-cu12==12.4.127",
    "triton==3.1.0",
]

# 工具类依赖
UTILS_REQUIREMENTS = [
    "numpy==2.0.2",
    "pandas==2.2.3",
    "pillow==10.2.0",
    "pyarrow==18.0.0",
    "PyYAML==6.0.2",
    "tqdm==4.67.0",
]

# 开发环境依赖
DEV_REQUIREMENTS = [
    "ipykernel",
    "ipython",
    "jupyter_client",
    "jupyter_core",
]

setup(
    name="my_training_project",
    version="0.1.0",
    description="Training project with PyTorch",
    author="Your Name",
    author_email="your.email@example.com",
    packages=find_packages(),
    python_requires=">=3.9",
    install_requires=[
        *MAIN_REQUIREMENTS,
        *UTILS_REQUIREMENTS,
        "aiohappyeyeballs==2.4.3",
        "aiohttp==3.11.0",
        "filelock==3.16.1",
        "fsspec==2024.9.0",
        "huggingface-hub==0.26.2",
        "tokenizers==0.20.3",
        "safetensors==0.4.5",
    ],
    extras_require={
        "cuda": CUDA_REQUIREMENTS,
        "dev": DEV_REQUIREMENTS,
    },
    dependency_links=[
        "https://download.pytorch.org/whl/cu121/torch/",
        "https://download.pytorch.org/whl/cu121/torchvision/",
        "https://download.pytorch.org/whl/cu121/torchaudio/",
    ]
)